import torch
import torch.nn as nn
from torchvision import transforms
from PIL import Image
import numpy as np
import faiss
import os

class FeatureExtractor:
    def __init__(self):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = self._load_model()
        self.preprocess = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225]
            )
        ])

    def _load_model(self):
        model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)
        model = nn.Sequential(*list(model.children())[:-1])
        return model.to(self.device).eval()

    def extract(self, image_path):
        try:
            img = Image.open(image_path).convert('RGB')
            img_tensor = self.preprocess(img).unsqueeze(0).to(self.device)
            with torch.no_grad():
                features = self.model(img_tensor).flatten()
            features /= torch.norm(features, p=2)  # L2归一化
            return features.cpu().numpy().astype('float32')
        except Exception as e:
            print(f"Error processing {image_path}: {str(e)}")
            return None